from torch import nn


class Loss_yolov1(nn.Module):
    def __init__(self, class_num):
        super(Loss_yolov1, self).__init__()
        self.class_num = class_num
    def forward(self, pred, labels):
        """

        :param pred:
        :param labels:
        :return:
        """

        num_gridx, num_gridy = labels.size()[-2:]     # 7*7  or 13*13
        num_b = 2
        num_cls = self.class_num
        noobj_confi_loss = 0.  # 不含目标的网格损失(只有置信度的损失)
        coor_loss = 0.  # 含有目标的bbox的坐标(coordinate)损失
        obj_confi_loss = 0.  # 含有目标的bbox的置信度损失
        class_loss = 0.  # 含有目标的网格的类别损失
        n_batch = labels.size()[0]

        for i in range(n_batch):
            for n in range(num_gridx):
                for m in range(num_gridy):
                    if labels[i, 4, m, n] == 1:  # (x, y, w, h, class_id)
                        # 将数据(x, y, w, h) 转成(x1, y1, x2, y2)
                        #   1. 先将px,py转换为cx,cy，即相对网格的位置转换为标准化后实际的bbox中心位置cx,xy
                        #   2. 然后再利用(cx-w/2,cy-h/2,cx+w/2,cy+h/2)转换为xyxy形式，用于计算iou
                        bbox1_pred_xyxy = ((pred[i, 0, m, n] + m) / num_gridx - pred[i, 2, m, n] / 2,
                                           (pred[i, 1, m, n] + n) / num_gridy - pred[i, 3, m, n] / 2,
                                           (pred[i, 0, m, n] + m) / num_gridx + pred[i, 2, m, n] / 2,
                                           (pred[i, 1, m, n] + n) / num_gridy + pred[i, 3, m, n] / 2)

        return None
